import argparse
import json
from typing import List

from pydantic import BaseModel
from tqdm import tqdm
from transformers import AutoTokenizer
from vllm import LLM, SamplingParams

parser = argparse.ArgumentParser()
parser.add_argument("-i", "--input-file-path", type=str)
parser.add_argument("-o", "--output-file-path", type=str)
args = parser.parse_args()

# 读取文件内容
with open(args.input_file_path, "r", encoding="utf-8") as file:
    sentences = [
        json.loads(line.strip())["CONCATENATED_CONTENT"]
        for line in file
        if line.strip()
    ]

# filter out empty sentences
sentences = [sentence for sentence in sentences if sentence]

# 按批次处理数据，每批次100条
batch_size = 1000
batches = [sentences[i : i + batch_size] for i in range(0, len(sentences), batch_size)]

# 模型和tokenizer初始化
model_path = "Qwen/Qwen2-7B-Instruct"  # 模型路径
tokenizer = AutoTokenizer.from_pretrained(model_path)
sampling_params = SamplingParams(temperature=0, top_p=0.9, max_tokens=512)
prefix_cached_llm = LLM(
    model=model_path,
    enable_prefix_caching=True,
    disable_custom_all_reduce=True,
    gpu_memory_utilization=0.9,
    tensor_parallel_size=1,
    max_model_len=4096,
)


class Response(BaseModel):
    raw: str = ""
    complaint: str = ""
    keywords: List[str]

    @classmethod
    def from_str(cls, s: str, raw: str):
        start = s.find("{")
        end = s.rfind("}") + 1
        try:
            obj = cls.model_validate_json(s[start:end])
        except Exception:
            obj = Response(keywords=[])
        obj.raw = raw
        return obj


system_prompt = """
你是广州智能政务话务员，专门做市民的投诉回访工作，你能够通过对话记录分析市民投诉的原因，并准确返回相关投诉原因的关键词，以用于聚类分析。
你的能力有:
- 分析对话内容，了解市民投诉的原因。
- 提取市民投诉的关键词，以便后续聚类分析。
- 如果电话没有接通，或者市民没有回话，你可以返回空列表。

你需要以json的格式返回关键词，例如:
{
    "complaint": "投诉原因xxx",
    "keywords": ["xxx", "xxx", "xxx", ...],
}
麻烦你先详细分析情况，然后再生成json结果，complaint字段中的内容再描述清楚一点。
"""

example_user_prompt_1 = """
您好，这边是广州市12345政府服务中心会员工号7716接到您主要是针对之前反映的这个噪音的问题进行回访，那这个情况是哪一方面不满意呢，一是问题没有得到解决方时承办单位回复的情况与实际情况不一致的。 你好。 呃，他你 们的人是我你们的人就是街道那边说派人去看了，但是呢，你派人去看我我我反映的是六点半到七点钟它那里是有有噪音的，他说他九点钟上班去看那 人家都结束了，那看什么了呢。 那我也知道我九点钟才上班，我也知道我九点钟才去看呢。 啊，就是处理效率不满意，但是这个噪音现在已经不存在 了，对吧。 我不知道赵颖还存不存在，因为他不是每天都吵他就是偶尔那么一两天吵。 偶尔一两天吵好的明白的那这边会记录您不满意情况，那感谢 您的接听，再见。
"""

example_assistant_prompt_1 = """
根据提供的对话记录，市民投诉的原因主要集中在以下几点：

1. 噪音问题未能在实际发生时间得到有效处理。
2. 承办单位的处理时间与市民反映的实际情况不一致。
3. 对处理效率不满意。

以下是详细的分析和提取的关键词：

```json
{
    "complaint": "市民反映的噪音问题在早上六点半到七点钟之间发生，但承办单位在九点钟上班后才去查看，导致问题未能在实际发生时间得到处理。市民对处理效率不满意，并认为承办单位的处理时间与实际情况不一致。",
    "keywords": ["噪音", "处理时间", "效率", "承办单位", "不一致"]
}
```
"""

example_user_prompt_2 = """您好，这边是广州市12345端口服务中心会员工号7416，见到您主要是针对之前反映的这个朱村街道的一个行政作风问题进行回访工作的那请问您当时诉求的情况，现在解决了没有呢。 你好。 您好，请问能听得清楚吗。 您好，请问能听得清楚吗。 您好，请问能听得清楚吗，这边电话已经接通。 您好，针对您之前反映的这个周春街道的问题进行回访，工作很抱歉，听不到有效的回 应将进行挂机的处理。 您好，很抱歉，听不到您有效的回应将进行挂机的处理。 您好，很抱歉，听不到有效的回应将进行挂机的处理，感谢您的接听 ，再见。"""

example_assistant_prompt_2 = """
根据提供的对话记录，由于市民未能给予有效回应，无法分析出具体的投诉原因和关键词。

以下是生成的json结果：

```json
{
    "complaint": "",
    "keywords": []
}
```
请提供其他对话记录以进行进一步的分析。
"""

example_user_prompt_3 = """您好，这边是广州市12345政府服务热线会员工号7416今当您主要是针对之前反映的这个花都交警大队的这个问题进行回访工作的那当时您反映他们的一个行政作风的问题的情况，现在解决了没有呢。 没解决，还是依照原发。 嗯，那 他们有没有针对您这个诉求回复过您处理的意见呢呢，用短信或者电话的一个回复内容的。 没有没有。 都没有好的这边已经记录那因为你给他们这个 处理意见进行满意度的评分一分是最低分五分是最高分，您是如何评价呢。 你们公司最低分吗。 一分是最低分也可以评零分的。 嗯，是评价花都交警还是评价那个政府的服务人员。 花都呃，您当时这个是花都交警大队的这个行政作风问题。 一分一分。 一分好的这边已经记录您不满意的情况，那感谢您的接听，再见，不客气。 好的。
"""

example_assistant_prompt_3 = """
根据提供的对话记录，市民投诉的原因主要集中在以下几点：

1. 花都交警大队的行政作风问题未得到解决。
2. 市民未收到任何关于投诉处理意见的回复（短信或电话）。
3. 市民对花都交警大队的处理非常不满意，给予最低评分。

以下是详细的分析和提取的关键词：

```json
{
    "complaint": "市民反映花都交警大队的行政作风问题没有得到解决，并且未收到任何关于投诉处理意见的回复（短信或电话）。市民对处理非常不满意，给予最低评分。",
    "keywords": ["花都交警大队", "行政作风", "未解决", "无回复", "不满意"]
}
```
"""


# 初始化结果列表
all_results = []


with open(args.output_file_path, "a", encoding="utf-8") as file:
    # 逐批处理数据
    for batch_index, batch in tqdm(enumerate(batches)):
        generating_prompts = []
        for prompt in batch:
            messages = [
                {"role": "system", "content": system_prompt},
                {"role": "user", "content": example_user_prompt_1},
                {"role": "assistant", "content": example_assistant_prompt_1},
                {"role": "user", "content": example_user_prompt_2},
                {"role": "assistant", "content": example_assistant_prompt_2},
                {"role": "user", "content": example_user_prompt_3},
                {"role": "assistant", "content": example_assistant_prompt_3},
                {"role": "user", "content": prompt},
            ]
            text = tokenizer.apply_chat_template(
                messages, tokenize=False, add_generation_prompt=True
            )
            generating_prompts.append(text)

        # Warmup so that the shared prompt's KV cache is computed.
        if batch_index == 0:
            prefix_cached_llm.generate(generating_prompts[0], sampling_params)

        # Generate with prefix caching.
        outputs = prefix_cached_llm.generate(generating_prompts, sampling_params)

        # Collect results
        batch_results: List[Response] = []
        for history, output in zip(batch, outputs):
            generated_text = output.outputs[0].text
            response = Response.from_str(generated_text, raw=history)
            batch_results.append(response)

        # 将每个批次的结果追加到输出文件中
        for result in batch_results:
            file.write(result.model_dump_json())
            file.write("\n")
            file.flush()
